package com.example.cs4.web;

import com.example.cs4.mapper.ChatHistoryMapper;
import com.example.cs4.model.ChatHistory;
import com.example.cs4.service.ChatHistoryService;
import com.example.cs4.service.MessageBroadcaster;
import com.example.cs4.service.MessageRankService;
import com.example.cs4.service.OnlineUserService;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.env.Environment;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@Component
public class SimpleChatHandler extends TextWebSocketHandler {

    @Autowired
    private ChatHistoryService chatHistoryService;

    @Autowired
    private MessageBroadcaster messageBroadcaster;

    @Autowired
    private OnlineUserService onlineUserService;

    @Autowired
    private MessageRankService messageRankService;

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    private final Set<WebSocketSession> sessions = ConcurrentHashMap.newKeySet();
    private final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();
    private final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Autowired
    private Environment env;

    @Autowired
    private ChatHistoryMapper chatHistoryMapper;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        session.setTextMessageSizeLimit(5000 * 1024);
        session.setBinaryMessageSizeLimit(10 * 1024 * 1024);

        String username = extractUsernameFromUri(session.getUri().getQuery());
        if (username == null || username.isEmpty()) {
            username = "用户_" + Long.toString(System.currentTimeMillis(), 36);
        }
        if (username.length() > 20) {
            username = username.substring(0, 20);
        }

        if (userSessions.containsKey(username)) {
            System.out.println("用户已在线: " + session.getId() + " 用户名: " + username);
            return;
        }

        sessions.add(session);
        userSessions.put(username, session);
        sessionMap.put(session.getId(), session);
        System.out.println("新客户端连接: " + session.getId() + " 用户名: " + username);
        sendUserListToAll();
        messageBroadcaster.broadcast(Collections.singletonMap("type", "JOIN"));
        messageBroadcaster.broadcast(Collections.singletonMap("user", username));

        onlineUserService.userLogin(username);

        messageBroadcaster.addSession(session);

        sendTopUsersToClient(session);

        sendUnreadCountToAll();
    }

    private void sendTopUsersToClient(WebSocketSession session) throws IOException {
        Set<ZSetOperations.TypedTuple<String>> topUsers = messageRankService.getTopUsers();
        List<Map<String, Object>> topUserList = topUsers.stream()
                .map(tuple -> {
                    Map<String, Object> userRank = new HashMap<>();
                    userRank.put("username", tuple.getValue());
                    userRank.put("messageCount", tuple.getScore());
                    return userRank;
                })
                .collect(Collectors.toList());

        Map<String, Object> rankMessage = new HashMap<>();
        rankMessage.put("type", "RANK");
        rankMessage.put("topUsers", topUserList);

        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(rankMessage)));
    }

    private String extractUsernameFromUri(String query) {
        if (query != null) {
            String[] params = query.split("&");
            for (String param : params) {
                String[] keyValue = param.split("=");
                if (keyValue.length > 1 && "user".equals(keyValue[0])) {
                    return keyValue[1];
                }
            }
        }
        return null;
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        Map<String, String> msg = objectMapper.readValue(message.getPayload(), Map.class);
        ensureMessageFields(msg);
        String type = msg.get("type");
        String user = msg.get("user");
        String messageContent = msg.get("message");
        String filename = msg.get("filename");
        String voice = msg.get("voice");
        String voiceType = msg.get("voiceType");
        Boolean isRead = msg.get("isRead") != null && msg.get("isRead").equals("true");
        String senderUsername = userSessions.entrySet().stream()
                .filter(entry -> entry.getValue().equals(session))
                .map(Map.Entry::getKey)
                .findFirst()
                .orElse("未知用户");

        if (!isMessageRateAllowed(senderUsername)) {
            Map<String, String> errorMessage = new HashMap<>();
            errorMessage.put("type", "ERROR");
            errorMessage.put("text", "消息发送频率过高，请稍后再试");
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
            return;
        }

        if (messageContent == null || messageContent.isEmpty()) {
            messageContent = "无内容";
            msg.put("message", messageContent);
        }

        if (type == null || "undefined".equals(type)) {
            Map<String, String> errorMessage = new HashMap<>();
            errorMessage.put("type", "ERROR");
            errorMessage.put("text", "未知消息类型: " + type);
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
            return;
        }

        else if ("GET_HISTORY".equals(type)) {
            if ("UNKNOWN".equals(user)) {
                List<ChatHistory> history = chatHistoryService.getAllChatHistory();
                Map<String, Object> historyResponse = new HashMap<>();
                historyResponse.put("type", "HISTORY");
                historyResponse.put("messages", history);
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(historyResponse)));
            } else {
                List<ChatHistory> history = chatHistoryService.getChatHistoryByReceiverOrSender(senderUsername, user);
                List<Integer> messageIds = history.stream().map(ChatHistory::getId).collect(Collectors.toList());
                chatHistoryService.markMessagesAsReadByIds(messageIds);
                Map<String, Object> historyResponse = new HashMap<>();
                historyResponse.put("type", "HISTORY");
                historyResponse.put("messages", history);
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(historyResponse)));

                int unreadCount = chatHistoryService.getUnreadMessageCount(senderUsername);
                Map<String, Object> unreadCountMessage = new HashMap<>();
                unreadCountMessage.put("type", "UNREAD_COUNT");
                unreadCountMessage.put("count", unreadCount);
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(unreadCountMessage)));
                return;
            }
        }

        else if ("GET_RANK".equals(type)) {
            Set<ZSetOperations.TypedTuple<String>> topUsers = messageRankService.getTopUsers();
            List<Map<String, Object>> topUserList = topUsers.stream()
                    .map(tuple -> {
                        Map<String, Object> userRank = new HashMap<>();
                        userRank.put("username", tuple.getValue());
                        userRank.put("messageCount", tuple.getScore());
                        return userRank;
                    })
                    .collect(Collectors.toList());

            Map<String, Object> rankMessage = new HashMap<>();
            rankMessage.put("type", "RANK");
            rankMessage.put("topUsers", topUserList);

            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(rankMessage)));
        }

        else if ("all".equals(user)) {
            msg.put("sender", senderUsername);
            msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream"));
            msg.put("voiceType", voiceType);
            System.out.println("Broadcasting public message: " + msg);
            broadcastMessageExcludingSender(msg, session);
            saveMessageToDatabase(senderUsername, "all", messageContent, msg.get("timestamp"), msg.get("fileType"), msg.get("file"), filename, voice, voiceType);
            updateMessageRank(senderUsername);
            sendUnreadCountToAll();
        } else if ("FILE".equals(type)) {
            try {
                msg.put("sender", senderUsername);
                msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream"));

                if ("all".equals(user)) {
                    System.out.println("Broadcasting public file message: " + msg);
                    messageBroadcaster.broadcast(msg);
                    saveMessageToDatabase(senderUsername, "all", messageContent, msg.get("timestamp"), msg.get("fileType"), msg.get("file"), filename, voice, voiceType);
                    messageRankService.incrementMessageCount(senderUsername);
                    sendUnreadCountToAll();
                } else {
                    WebSocketSession targetSession = userSessions.get(user);
                    if (targetSession == null) {
                        targetSession = sessionMap.get(user);
                    }
                    if (targetSession != null && targetSession.isOpen()) {
                        msg.put("sender", senderUsername);
                        msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream"));
                        targetSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                        if (!targetSession.equals(session)) {
                            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                        }
                        saveMessageToDatabase(senderUsername, user, messageContent, msg.get("timestamp"), msg.get("fileType"), msg.get("file"), filename, voice, voiceType);
                        chatHistoryService.markMessagesAsReadByIds(Collections.singletonList(chatHistoryService.getChatHistoryByReceiverOrSender(user, senderUsername).get(0).getId()));
                        sendUnreadCountToAll();
                    } else {
                        System.out.println("目标用户不在线或不存在: " + user);
                        Map<String, String> errorMessage = new HashMap<>(msg);
                        errorMessage.put("type", "ERROR");
                        errorMessage.put("text", "目标用户不在线或不存在: " + user);
                        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
                    }
                }
            } catch (Exception e) {
                System.err.println("Error handling file message: " + e.getMessage());
                e.printStackTrace();
                Map<String, String> errorMessage = new HashMap<>();
                errorMessage.put("type", "ERROR");
                errorMessage.put("text", "文件发送失败: " + e.getMessage());
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
            }
        } else if ("VOICE".equals(type)) {
            try {
                msg.put("sender", senderUsername);
                msg.put("voiceType", voiceType);

                if ("all".equals(user)) {
                    System.out.println("Broadcasting public voice message: " + msg);
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                    broadcastMessage(msg);
                    saveMessageToDatabase(senderUsername, "all", messageContent, msg.get("timestamp"), null, null, null, voice, voiceType);
                    messageRankService.incrementMessageCount(senderUsername);
                    sendUnreadCountToAll();
                } else {
                    WebSocketSession targetSession = userSessions.get(user);
                    if (targetSession == null) {
                        targetSession = sessionMap.get(user);
                    }
                    if (targetSession != null && targetSession.isOpen()) {
                        targetSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                        saveMessageToDatabase(senderUsername, user, messageContent, msg.get("timestamp"), null, null, null, voice, voiceType);
                        chatHistoryService.markMessagesAsReadByIds(Collections.singletonList(chatHistoryService.getChatHistoryByReceiverOrSender(user, senderUsername).get(0).getId()));
                        sendUnreadCountToAll();
                    } else {
                        System.out.println("目标用户不在线或不存在: " + user);
                        Map<String, String> errorMessage = new HashMap<>(msg);
                        errorMessage.put("type", "ERROR");
                        errorMessage.put("text", "目标用户不在线或不存在: " + user);
                        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
                    }
                }
            } catch (Exception e) {
                System.err.println("Error handling voice message: " + e.getMessage());
                e.printStackTrace();
                Map<String, String> errorMessage = new HashMap<>();
                errorMessage.put("type", "ERROR");
                errorMessage.put("text", "语音发送失败: " + e.getMessage());
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
            }
        } else {
            if (senderUsername.equals(user)) {
                Map<String, String> errorMessage = new HashMap<>();
                errorMessage.put("type", "ERROR");
                errorMessage.put("text", "不能给自己发送私聊消息");
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
                return;
            }

            WebSocketSession targetSession = userSessions.get(user);
            if (targetSession == null) {
                targetSession = sessionMap.get(user);
            }
            if (targetSession != null && targetSession.isOpen()) {
                msg.put("sender", senderUsername);
                msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream"));
                msg.put("voiceType", voiceType);

                targetSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                saveMessageToDatabase(senderUsername, user, messageContent, msg.get("timestamp"), msg.get("fileType"), msg.get("file"), filename, voice, voiceType);
                chatHistoryService.markMessagesAsReadByIds(Collections.singletonList(chatHistoryService.getChatHistoryByReceiverOrSender(user, senderUsername).get(0).getId()));
                sendUnreadCountToAll();
            } else {
                System.out.println("目标用户不在线或不存在: " + user);
                Map<String, String> errorMessage = new HashMap<>(msg);
                errorMessage.put("type", "ERROR");
                errorMessage.put("text", "目标用户不在线或不存在: " + user);
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
            }
        }
    }

    private void saveMessageToDatabase(String sender, String receiver, String message, String timestamp, String fileType, String file, String filename, String voice, String voiceType) throws JsonProcessingException {
        ChatHistory chatHistory = new ChatHistory();
        chatHistory.setSender(sender);
        chatHistory.setReceiver(receiver);
        chatHistory.setMessage(message);
        chatHistory.setFileType(fileType);
        chatHistory.setFile(file);
        chatHistory.setFilename(filename);
        chatHistory.setVoice(voice);
        chatHistory.setVoiceType(voiceType);
        chatHistory.setIsRead(false);

        if (voice != null && !voice.isEmpty() && (chatHistory.getVoiceFilename() == null || chatHistory.getVoiceFilename().isEmpty())) {
            chatHistory.setVoiceFilename("voice_" + System.currentTimeMillis() + ".wav");
        }

        try {
            chatHistory.setTimestamp(new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse(timestamp));
        } catch (Exception e) {
            chatHistory.setTimestamp(new Date());
        }
        chatHistoryService.saveChatHistory(chatHistory);

        String messageKey = "chat:message:" + chatHistory.getId();
        redisTemplate.opsForValue().set(messageKey, objectMapper.writeValueAsString(chatHistory));
        redisTemplate.expire(messageKey, 7, TimeUnit.DAYS);
    }

    private void ensureMessageFields(Map<String, String> msg) {
        if (!msg.containsKey("sender")) {
            msg.put("sender", "未知用户");
        }
        if (!msg.containsKey("timestamp")) {
            msg.put("timestamp", new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").format(new Date()));
        }
        if ("FILE".equals(msg.get("type")) && !msg.containsKey("fileType")) {
            msg.put("fileType", "application/octet-stream");
        }
        if (!msg.containsKey("file")) {
            msg.put("file", "");
        }
        if (!msg.containsKey("filename")) {
            msg.put("filename", "");
        }
        if (!msg.containsKey("voiceFilename")) {
            msg.put("voiceFilename", "");
        }
        if (!msg.containsKey("type")) {
            msg.put("type", "UNKNOWN");
        }
        if (!msg.containsKey("user")) {
            msg.put("user", "UNKNOWN");
        }
        if (!msg.containsKey("message")) {
            msg.put("message", "");
        }
        if (!msg.containsKey("filename")) {
            msg.put("filename", "");
        }
        if (!msg.containsKey("voice")) {
            msg.put("voice", "");
        }
        if (!msg.containsKey("voiceType")) {
            msg.put("voiceType", "");
        }
        if (msg.containsKey("voice") && !msg.get("voice").isEmpty() && (msg.get("voiceFilename") == null || msg.get("voiceFilename").isEmpty())) {
            msg.put("voiceFilename", "voice_" + System.currentTimeMillis() + ".wav");
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String username = userSessions.entrySet().stream()
                .filter(entry -> entry.getValue().equals(session))
                .map(Map.Entry::getKey)
                .findFirst()
                .orElse("未知用户");
        userSessions.values().remove(session);
        sessionMap.remove(session.getId());
        System.out.println("客户端断开连接: " + session.getId());
        sendUserListToAll();
        messageBroadcaster.broadcast(Collections.singletonMap("type", "LEAVE"));
        messageBroadcaster.broadcast(Collections.singletonMap("user", username));

        onlineUserService.userLogout(username);

        messageBroadcaster.removeSession(session);
    }

    private void broadcastMessage(Map<String, String> message, Map<String, String> typeInfo) {
        try {
            for (WebSocketSession session : sessions) {
                if (session.isOpen()) {
                    Map<String, String> combinedMessage = new HashMap<>(message);
                    combinedMessage.putAll(typeInfo);
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(combinedMessage)));
                }
            }
        } catch (IOException e) {
            System.err.println("Error broadcasting message to session");
            e.printStackTrace();
        }
    }

    private void broadcastMessageExcludingSender(Map<String, String> message, WebSocketSession senderSession) {
        try {
            for (WebSocketSession session : sessions) {
                if (session.isOpen() && !session.equals(senderSession)) {
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
                }
            }
        } catch (IOException e) {
            System.err.println("Error broadcasting message to session");
            e.printStackTrace();
        }
    }

    private void broadcastMessage(Map<String, String> message) {
        try {
            for (WebSocketSession session : sessions) {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
                }
            }
        } catch (IOException e) {
            System.err.println("Error broadcasting message to session");
            e.printStackTrace();
        }
    }

    private void sendUserListToAll() {
        Set<String> users = userSessions.keySet();
        try {
            for (WebSocketSession session : sessions) {
                if (session.isOpen()) {
                    Map<String, Object> userListMessage = new HashMap<>();
                    userListMessage.put("type", "USERS");
                    userListMessage.put("users", users);
                    String message = objectMapper.writeValueAsString(userListMessage);
                    session.sendMessage(new TextMessage(message));
                }
            }
        } catch (IOException e) {
            System.err.println("Error sending user list");
            e.printStackTrace();
        }
    }

    private List<ChatHistory> getAllChatHistory() {
        return chatHistoryService.getAllChatHistory();
    }

    private void updateMessageRank(String username) {
        messageRankService.incrementMessageCount(username);
    }


    private void markMessagesAsRead(List<Integer> messageIds) throws IOException {
        chatHistoryMapper.markMessagesAsReadByIds(messageIds);
        sendUnreadCountToAll();
    }

    private boolean isMessageRateAllowed(String userId) {
        String serverPort = env.getProperty("server.port");
        String rateKey = "chat:message_rate:" + userId + ":" + serverPort;
        Long count = redisTemplate.opsForValue().increment(rateKey);
        if (count == 1) {
            redisTemplate.expire(rateKey, 1, TimeUnit.MINUTES);
        }
        return count <= 10;
    }

    private void sendUnreadCountToAll() throws IOException {
        for (WebSocketSession session : sessions) {
            if (session.isOpen()) {
                String username = userSessions.entrySet().stream()
                        .filter(entry -> entry.getValue().equals(session))
                        .map(Map.Entry::getKey)
                        .findFirst()
                        .orElse("未知用户");
                int unreadCount = chatHistoryService.getUnreadMessageCount(username);
                if (unreadCount > 0) {
                    Map<String, Object> unreadCountMessage = new HashMap<>();
                    unreadCountMessage.put("type", "UNREAD_COUNT");
                    unreadCountMessage.put("count", unreadCount);
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(unreadCountMessage)));
                }
            }
        }
    }
}